import pandas as pd
import numpy as np
from xgboost import XGBRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

# 读取数据
df = pd.read_excel(r'stockdata.xlsx')

# 转换日期为datetime格式
df['日期'] = pd.to_datetime(df['日期'])

# 提取年份
df['年份'] = df['日期'].dt.year

df['股票代码'] = df['股票代码'].astype(str)
df['行业代码'] = df['行业代码'].astype(str)

# 填充缺失值
numeric_columns = df.select_dtypes(include=[np.number]).columns
df[numeric_columns] = df[numeric_columns].fillna(df[numeric_columns].mean())

# 提取相应的特征和标签
features = ['股票代码', '年份', '实际控制人性质', '行业代码', '短期负债', '长期负债合计']
label = '资产负债率'

# 构建特征矩阵和标签向量
X = df[features]
y = df[label]

# 将非数值特征转换为数值
X = pd.get_dummies(X, columns=['股票代码', '实际控制人性质', '行业代码'])

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建并训练XGBoost模型
model = XGBRegressor(n_estimators=100, learning_rate=0.1, max_depth=3, random_state=42)
model.fit(X_train, y_train)

# 评估模型
y_pred = model.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
print(f'Test MSE: {mse}')

# 预测未来10年的资产负债率
future_years = np.arange(df['年份'].max() + 1, df['年份'].max() + 11)
future_data = []

for year in future_years:
    for _, row in df.iterrows():
        future_data.append([row['股票代码'], year, row['实际控制人性质'], row['行业代码'], row['短期负债'], row['长期负债合计']])

future_df = pd.DataFrame(future_data, columns=features)
future_df = pd.get_dummies(future_df, columns=['股票代码', '实际控制人性质', '行业代码'])

# 确保未来数据的特征与训练数据的特征对齐
future_df = future_df.reindex(columns=X_train.columns, fill_value=0)

future_predictions = model.predict(future_df)

# 将预测结果保存到 DataFrame
result_df = future_df[['股票代码', '年份', '实际控制人性质', '行业代码']]
result_df['预测资产负债率'] = future_predictions

# 保存结果到Excel
result_df.to_excel('future_debt_ratio_predictions.xlsx', index=False)